#include <stdio.h>

#define N 1024
#define TPB 256  // 每个线程块的线程数

__global__ void add(int *a, int *b, int *c) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < N) {
        c[tid] = a[tid] + b[tid];
    }
}

int main() {
    int *a, *b, *c;        // 主机端指针
    int *d_a, *d_b, *d_c;  // 设备端指针
    int size = N * sizeof(int);

    // 分配主机内存
    a = (int*)malloc(size);
    b = (int*)malloc(size);
    c = (int*)malloc(size);

    // 初始化输入数组
    for (int i = 0; i < N; i++) {
        a[i] = i;
        b[i] = i * 2;
    }

    // 分配设备内存
    cudaMalloc(&d_a, size);
    cudaMalloc(&d_b, size);
    cudaMalloc(&d_c, size);

    // 将数据从主机复制到设备
    cudaMemcpy(d_a, a, size, cudaMemcpyHostToDevice);
    cudaMemcpy(d_b, b, size, cudaMemcpyHostToDevice);

    // 计算启动的线程块数量
    int blocks = (N + TPB - 1) / TPB;

    // 启动核函数
    add<<<blocks, TPB>>>(d_a, d_b, d_c);

    // 将结果从设备复制回主机
    cudaMemcpy(c, d_c, size, cudaMemcpyDeviceToHost);

    // 验证结果
    printf("验证结果中...\n");
    bool success = true;
    for (int i = 0; i < N; i++) {
        if (c[i] != a[i] + b[i]) {
            printf("验证失败：c[%d] = %d, 应该是 %d\n", i, c[i], a[i] + b[i]);
            success = false;
            break;
        }
    }
    if (success) {
        printf("验证成功！所有 %d 个元素都计算正确。\n", N);
    }

    // 释放设备内存
    cudaFree(d_a);
    cudaFree(d_b);
    cudaFree(d_c);

    // 释放主机内存
    free(a);
    free(b);
    free(c);

    return 0;
}